import os
import base64
from google import genai
from google.genai import types
from google.cloud import storage
from google.genai.types import HttpOptions, Part

class GoogleSpeechTranscriber:
    def __init__(self, project_id, location="us-central1", model="gemini-2.0-flash-001", bucket="cloud-ai-platform-1a06e8cb-b424-44c0-a95e-c28718a96ba3"):
        self.client = genai.Client(vertexai=True, project=project_id, location=location)
        self.model = model
        self.bucket = bucket

    def upload_to_gcs(self, source_file_path):
        print("Uploading to GCS...")
        storage_client = storage.Client()
        bucket = storage_client.bucket(self.bucket)
        filename = os.path.basename(source_file_path)
        destination_blob_name = f"audio_uploads/{filename}"
        blob = bucket.blob(destination_blob_name)
        blob.upload_from_filename(source_file_path)
        gcs_uri = f"gs://{self.bucket}/{destination_blob_name}"
        print(f"Uploaded to {gcs_uri}")
        return gcs_uri

    def delete_from_gcs(self, source_file_path):
        print("Deleting GCS file...")
        filename = os.path.basename(source_file_path)
        destination_blob_name = f"audio_uploads/{filename}"
        storage_client = storage.Client()
        bucket = storage_client.bucket(self.bucket)
        blob = bucket.blob(destination_blob_name)
        blob.delete()
        print(f"Deleted gs://{self.bucket}/{destination_blob_name}")

    def transcribe_audio(self, audio_file_path):
        if not os.path.exists(audio_file_path):
            return {"error": f"File not found: {audio_file_path}"}

        print(f"音频文件: {audio_file_path}")
        # gcs_uri = self.upload_to_gcs(audio_file_path)
        with open(audio_file_path, "rb") as f:
            audio_bytes = f.read()

        base64_audio_str = base64.b64encode(audio_bytes).decode("utf-8")  # 字符串形式
        print("✅ Base64 字符串（可粘贴到代码中）如下：")
        print(base64_audio_str[:100] + "...")  # 打印前100个字符看看

        audio_part = types.Part.from_bytes(
            data=audio_bytes,
            mime_type="audio/mpeg",  # 或根据实际情况使用 "audio/wav", "audio/webm"
        )

        contents = [
            types.Content(
                role="user",
                parts=[
                    # types.Part.from_uri(file_uri="gs://cloud-ai-platform-1a06e8cb-b424-44c0-a95e-c28718a96ba3/audio_uploads/41.mp3", mime_type="audio/mpeg"),
                    audio_part,
                    types.Part.from_text(
                        text="你收到的是一个音频内容，请识别出其中说的话及其语言，并将每句话翻译成中文。你是没有收到内容还是收到内容了无法正确识别为音频文件也告诉我一下"
                             "**输出格式要求：**\n"
                             "1. 检测到的源语言（如：马来语、泰语）。\n"
                             "2. 以 JSON 格式输出文本及翻译，时间戳精准到秒。\n"
                             "JSON 格式如下：\n"
                             "{\n"
                             '  "source_language": "马来语",\n'
                             '  "transcriptions": [\n'
                             "    {\n"
                             '      "start_time": "00s",\n'
                             '      "end_time": "05s",\n'
                             '      "original_text": "xxxx",\n'
                             '      "translated_text": "xxxx"\n'
                             "    }\n"
                             "  ]\n"
                             "}"
                    ),
                ],
            )
        ]

        generate_content_config = types.GenerateContentConfig(
            temperature=0.2,
            top_p=0.8,
            max_output_tokens=8192,
            response_modalities=["TEXT"],
            safety_settings=[
                types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="OFF"),
                types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="OFF"),
                types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="OFF"),
                types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="OFF"),
            ]
        )

        response_text = ""
        try:
            response = self.client.models.generate_content_stream(
                model=self.model,
                contents=contents,
                config=generate_content_config,
            )
            
            for chunk in response:
                print(chunk.text)  # 打印每个 chunk
                response_text += chunk.text
        except Exception as e:
            print(f"Error: {str(e)}")
            return {"error": str(e)}
        # finally:
            # self.delete_from_gcs(audio_file_path)

        return {"transcription": response_text}
